from contextvars import ContextVar
from typing import Dict, Any

from app.constants import MDC_TRACE_ID, SAME_TOKEN

# 创建一个上下文变量，用于存储线程局部数据
# ContextVar提供了比threading.local更安全的上下文管理方式，特别适用于异步环境
_thread_local: ContextVar[Dict[str, Any]] = ContextVar("context", default={})


# ======================== trace_id start ======================== #
def set_trace_id(trace_id: str) -> None:
    """
    设置当前上下文的trace_id

    :param trace_id: 链路追踪ID，用于日志跟踪和问题定位
    """
    _thread_local.get()[MDC_TRACE_ID] = trace_id


def get_trace_id() -> str:
    """
    获取当前上下文的trace_id

    :return: 当前的链路追踪ID，如果未设置则返回None
    """
    return _thread_local.get().get(MDC_TRACE_ID)


# ======================== trace_id end ======================== #

# ======================== user_id start ======================== #
def set_user_id(user_id: int) -> None:
    """
    设置当前上下文的用户ID

    :param user_id: 用户id
    """
    _thread_local.get()["user_id"] = user_id


def get_user_id() -> int:
    """
    获取当前上下文的用户ID

    :return: 当前用户ID，如果未设置则返回默认值0
    """
    return _thread_local.get().get("user_id", 0)


# ======================== user_id end ======================== #


# ======================== token start ======================== #
def set_token(token: str) -> None:
    """
    设置当前上下文的访问令牌

    :param token: 用户的访问令牌，用于身份验证
    """
    _thread_local.get()["token"] = token


def get_token() -> str:
    """
    获取当前上下文的访问令牌

    :return: 当前访问令牌，如果未设置则返回None
    """
    return _thread_local.get().get("token")


# ======================== token end ======================== #

# ======================== same_token start ======================== #
def set_same_token(same_token: str) -> None:
    """
    设置当前上下文的SAME_TOKEN

    :param same_token: same_token，用于微服务子服务外网隔离
    """
    _thread_local.get()[SAME_TOKEN] = same_token


def get_same_token() -> str:
    """
    获取当前上下文的SAME_TOKEN

    :return: same_token，如果未设置则返回None
    """
    return _thread_local.get().get(SAME_TOKEN)


# ======================== same_token end ======================== #

def remove() -> None:
    """ 清空当前上下文的所有数据 """
    _thread_local.set({})
